package com.boge.ai.rag.embedding;

import com.boge.ai.utils.LLMUtils;
import org.apache.commons.math3.linear.ArrayRealVector;
import org.apache.commons.math3.linear.RealVector;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class Demo3 {
    /**
     * 向量的相似度计算
     * @param args
     */
    public static void main(String[] args) {
        String query = "人工智能伦理";
        List<String> documents = new ArrayList<>();
        documents.add("全球科技巨头联合发布AI伦理指南，强调透明度与公平性");
        documents.add("欧盟通过新的法规，要求所有AI系统必须符合严格的隐私保护标准");
        documents.add("科学家警告：如果不加以控制，AI可能会加剧社会不平等");
        documents.add("国际会议讨论如何防止AI武器化，并呼吁制定国际条约");
        documents.add("某国政府宣布将投资数十亿美元用于支持可持续发展的AI技术研究");
        float[] embedding = LLMUtils.embedding(query, null);

        RealVector embeddingVector = new ArrayRealVector(LLMUtils.toDoubleArray(embedding));
        System.out.println("欧式距离");
        System.out.println(LLMUtils.l2(embeddingVector,embeddingVector));

        for (String document : documents) {
            float[] documentEmbedding = LLMUtils.embedding(document, null);


            double v = LLMUtils.l2(embeddingVector, new ArrayRealVector(LLMUtils.toDoubleArray(documentEmbedding)));
            System.out.println(v);
        }

        System.out.println("余旋");
        System.out.println(LLMUtils.cosSim(embeddingVector,embeddingVector));

        for (String document : documents) {
            float[] documentEmbedding = LLMUtils.embedding(document, null);


            double v = LLMUtils.cosSim(embeddingVector, new ArrayRealVector(LLMUtils.toDoubleArray(documentEmbedding)));
            System.out.println(v);
        }
    }
}
